# coding:utf-8

import cv2
from multiprocessing import Process
import numpy as np
import os
import random
import json

nakes_prefix = r"c:\Users\Administrator\Desktop\SSD\source"
target_prefix = r'c:\Users\Administrator\Desktop\SSD\target'
bg_dir = r"c:\Users\Administrator\Desktop\SSD\temp"
info_file = r"c:\Users\Administrator\Desktop\SSD\info.json"
sum_ = [0]
nakes_names = [nake for nake in os.listdir(nakes_prefix)]
category_map = dict(zip(sorted(nakes_names), range(0, len(nakes_names))))
nakes_dirs = [os.path.join(nakes_prefix, nake) for nake in nakes_names]
info = []
positions = []


def random_position(seed, bgshape):
    if positions:
        return
    h, w = bgshape
    i = 0
    while i < seed:
        x1 = random.randint(0, w)
        y1 = random.randint(0, h)
        if (x1, y1) not in positions:
            i += 1
            positions.append((x1, y1))


def build_dir(dirn):
    t = os.path.join(target_prefix, dirn)
    if not os.path.exists(t):
        os.makedirs(t)
    return t


def merge_pic(p_nake, p_bg, nake_dir):
    """
    :param p_nake: 实物图
    :param p_bg: 背景图
    :return:
    """
    bg = cv2.imread(p_bg)
    nake = cv2.imread(p_nake)
    nh, nw, channel = nake.shape
    h, w, channel = bg.shape
    if not positions:
        x1 = random.randint(0, w - nw - 1)
        y1 = random.randint(0, h - nh - 1)
    else:
        (x1, y1) = random.choice(positions)

    if y1 > h or x1 > w:
        print("x1 return")
        return
    x2 = x1 + nw
    y2 = y1 + nh
    if x2 > w or y2 > h:
        print(h, y2)
        print(w, x2)
        print("x2 return")
        return

    roi = bg[y1:y2, x1:x2]
    for i in range(0, nh):
        for j in range(0, nw):
            roi[i, j] = nake[i, j]

    bg[y1:y2, x1:x2] = roi

    nake_dir_name = os.path.split(nake_dir)[-1]
    t_dir = build_dir(nake_dir_name)
    t_name = os.path.split(p_nake)[-1].split(".")[0]
    t_bg = os.path.split(p_bg)[-1]
    t_name = t_name + "_" + t_bg
    cv2.imwrite(os.path.join(t_dir, t_name), bg)
    tg = {"file_name": os.path.join(t_dir, t_name),
          "width": w,
          "height": h,
          "annotations": [{"classes": nake_dir_name,
                           "xmin": x1,
                           "ymin": y1,
                           "xmax": x2,
                           "ymax": y2,
                           "category_id": category_map[nake_dir_name]}
                          ]
          }
    info.append(tg)
    sum_[0] += 1

def do(nake_dir):
    for nake_img in os.listdir(nake_dir):
        nake_img = os.path.join(nake_dir, nake_img)
        for bg in os.listdir(bg_dir):
            bg = os.path.join(bg_dir, bg)
            merge_pic(nake_img, bg, nake_dir)

def read():
    ps = []
    nakes_dirs=[os.path.join(nakes_prefix,i) for i in nakes_dir]
    for nake_dir in nakes_dirs:
        p = Process(target=do,args=(nake_dir,))
        ps.append(p)
    for p in ps:
        print("启动进程")
        p.start()
    for p in ps:
        p.join()


if __name__ == "__main__":
    print("start!")
    read()
    print("end!")
# coding:utf-8

def read():
    # bg = r"C:\Users\Administrator\Desktop\SSD\target\sfim.jpg"
    # fg = r'C:\Users\Administrator\Desktop\SSD\source\apple\6.jpg'
    # dirs = r'C:\Users\Administrator\Desktop\SSD'
    # merge_pic(fg,bg,dirs)
    random_position(2,(400,400))
    for nake_dir in nakes_dirs:
        do(nake_dir)
    import random
    random.shuffle(info)
    length = len(info)
    t_len = int(length*0.8)
    with open(train_file, 'w',encoding="utf-8") as f:
        infos = {"images":info[:t_len],
                "category":category_map}
        json.dump(infos,f)
    with open(val_file,'w',encoding="utf-8") as f:
        infos = {"images":info[t_len:],
                "category":category_map}
        json.dump(infos,f)